Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[MLX] [bugfix] Preserve dtype of array when converting to torch #1349

Merged
merged 3 commits into from
Jan 15, 2025
Merged

[MLX] [bugfix] Preserve dtype of array when converting to torch #1349

merged 3 commits into from
Jan 15, 2025

Conversation

neilmehta24
Copy link
Contributor

We have noticing the following error with a recent version of outlines when used with MLX:

TypeError: argument 'token_id': 'float' object cannot be interpreted as an integer

At:
  /.../outlines_core/fsm/guide.py(294): get_next_state
  /.../outlines/processors/structured.py(101): process_logits
  /.../outlines/processors/base_logits_processor.py(90): __call__

The issue is that the MLX array of tokens, which are integers, are being force-converted to floats, even though outlines expects an integer array. This is because all MLX arrays are being converted to float32, even when it's not necessarily appropriate, like in this case. Looking at the commented link, the advice was to convert to float32 only for bfloat16, because numpy does not support bfloat16. Now the MLX _to_torch implementation matches the other array libraries, none of the other libraries are being force-casted to float

@rlouf rlouf self-requested a review December 21, 2024 14:01
@neilmehta24
Copy link
Contributor Author

@rlouf The CI configuration prevents me from reaching the coverage requirements. From the CI logs:

outlines/processors/base_logits_processor.py (0.0%): Missing lines 113-115

But note that the CI doesn't install or run any MLX code, so to my understanding it's impossible to improve the CI coverage report when adding MLX-specific logic.

However, I did check locally that there isn't a regression in coverage, and that the tests are passing

Coverage:
Screenshot 2024-12-22 at 4 03 05 PM

Test:
Screenshot 2024-12-22 at 4 05 17 PM

So this should be good to go, please let me know what you think.

@neilmehta24
Copy link
Contributor Author

Hey @rlouf , what do you think of the comment above? Is this PR good to go?

@rlouf rlouf merged commit 088f439 into dottxt-ai:main Jan 15, 2025
5 of 6 checks passed
@rlouf
Copy link
Member

rlouf commented Jan 15, 2025

Yes, sorry for the delay in reviewing. Thank you for contributing this fix!

@neilmehta24 neilmehta24 deleted the mlx-to-torch-fix branch January 20, 2025 20:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@rlouf rlouf rlouf approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@neilmehta24 @rlouf